def overlapped_forward_backward(
    self,
    forward_chunk,  # the module of the forward chunk
    forward_inputs,
    forward_loss_fn_node,
    backward_chunk,  # the module of the backward chunk, maybe not used
    backward_loss_fn_node,
    backward_input_grads,
    scaler,
    p2p_async_handle=None,  # only used for p2p comm overlap
):          
    if backward_loss_fn_node is not None:
        if scaler:
            backward_input_grads = backward_loss_fn_node.backward(scaler=scaler)
        else:
            backward_input_grads = backward_loss_fn_node.backward()

    if p2p_async_handle is not None:
        p2p_async_handle.forward_handle_wait()
    forward_inputs = forward_chunk.forward(forward_inputs)
    
    if p2p_async_handle is not None:
        p2p_async_handle.forward_async_comm(forward_inputs)
        p2p_async_handle.backward_handle_wait()

    backward_input_grads = backward_chunk.backward(backward_input_grads)

    if p2p_async_handle is not None:
        p2p_async_handle.backwatd_async_comm(backward_input_grads)
            
    if forward_loss_fn_node is not None:
        forward_loss = forward_loss_fn_node.forward(forward_inputs)
    else:
        forward_loss = None

    return forward_inputs, forward_loss, backward_input_grads